use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

use tokio::sync::Mutex as TokioMutex;
use tracing::debug;

/// 具有日志记录和持有者信息的互斥锁。
///
/// `LoggedMutex<T>` 封装了 `tokio::sync::Mutex<T>`，并提供了额外的功能：
/// - 在获取和释放锁时记录日志。
/// - 跟踪当前持有锁的持有者的 ID，以及获取锁的位置（模块和行号）。
///
/// # 类型参数
///
/// * `T`: 互斥锁保护的数据的类型。
pub struct LoggedMutex<T> {
    inner: TokioMutex<T>,
    name: String,
    holder_id: Arc<AtomicUsize>,
    next_holder_id: AtomicUsize,
}

impl<T> LoggedMutex<T> {
    /// 创建一个新的 `LoggedMutex` 实例。
    ///
    /// # 参数
    ///
    /// * `name` - 互斥锁的名称，用于日志记录。
    /// * `value` - 互斥锁保护的初始值。
    pub fn new(name: &str, value: T) -> Self {
        LoggedMutex {
            inner: TokioMutex::new(value),
            name: name.to_string(),
            holder_id: Arc::new(AtomicUsize::new(0)), // 初始时没有持有者
            next_holder_id: AtomicUsize::new(1),
        }
    }

    /// 获取当前持有锁的持有者 ID。
    ///
    /// 如果锁当前未被任何线程持有，则返回 `None`。
    ///
    /// # 返回值
    ///
    /// * `Some(usize)`: 持有者的 ID。
    /// * `None`: 锁未被持有。
    pub fn current_holder(&self) -> Option<usize> {
        let holder_id = self.holder_id.load(Ordering::Relaxed);
        if holder_id == 0 {
            None
        } else {
            Some(holder_id)
        }
    }

    /// 尝试获取互斥锁。
    ///
    /// 如果锁立即可用，则此方法会立即返回一个 `LoggedMutexGuard`。
    /// 如果锁当前被另一个线程持有，则此方法会立即返回 `None`，而不会阻塞。
    ///
    /// # 返回值
    ///
    /// * `Some(LoggedMutexGuard)`: 成功获取锁。
    /// * `None`: 锁当前被其他线程持有。
    pub fn try_lock(&self) -> Option<LoggedMutexGuard<'_, T>> {
        match self.inner.try_lock() {
            Ok(guard) => {
                let holder_id = self.next_holder_id.fetch_add(1, Ordering::Relaxed);
                self.holder_id.store(holder_id, Ordering::Relaxed);
                debug!(
                    name = %self.name,
                    holder_id = holder_id,
                    "Lock (try_lock)"
                );
                Some(LoggedMutexGuard {
                    inner: guard,
                    mutex: self,
                    holder_id,
                })
            }
            Err(_) => {
                debug!(name = %self.name, "Try lock failed (try_lock)");
                None
            }
        }
    }

    /// 异步获取互斥锁。
    ///
    /// 此方法会返回一个 `Future`，它会在锁变为可用时解析为 `LoggedMutexGuard`。
    /// 如果锁当前被另一个线程持有，则此方法会挂起当前任务，直到锁被释放。
    ///
    /// # 返回值
    ///
    /// 一个 `Future`，它会解析为 `LoggedMutexGuard`。
    pub async fn lock(&self) -> LoggedMutexGuard<'_, T> {
        let guard = self.inner.lock().await;
        let holder_id = self.next_holder_id.fetch_add(1, Ordering::Relaxed);
        self.holder_id.store(holder_id, Ordering::Relaxed);
        debug!(
            name = %self.name,
            holder_id = holder_id,
            // module = module_path!(), // 添加模块名
            // line = line!(),             // 添加行号
            "Lock"
        );
        LoggedMutexGuard {
            inner: guard,
            mutex: self,
            holder_id,
        }
    }

    /// 获取对内部数据的可变原子引用。
    ///
    /// # 安全性
    ///
    /// 这个方法会绕过互斥锁的保护。在大多数情况下，你应该使用 `lock` 或 `try_lock` 方法来安全地访问受保护的数据。
    ///
    /// # 返回值
    ///
    /// 对受保护数据的可变引用。
    pub fn get_mut(&mut self) -> &mut T {
        self.inner.get_mut()
    }

    /// 执行一个闭包，该闭包持有锁。
    ///
    /// 这提供了一个方便的方法来在获取和释放锁的范围内执行代码。
    ///
    /// # 参数
    ///
    /// * `func` - 要执行的闭包。它接收对受保护数据的可变引用。
    ///
    /// # 类型参数
    ///
    /// * `F`: 闭包的类型。
    /// * `R`: 闭包的返回值类型。
    ///
    /// # 返回值
    ///
    /// 闭包的返回值。
    pub async fn execute<F, R>(&self, func: F) -> R
    where
        F: FnOnce(&mut T) -> R + Send,
        R: Send,
    {
        let mut guard = self.lock().await;
        func(&mut *guard)
    }
}

/// 互斥锁的持有 guard。
///
/// 当 `LoggedMutexGuard` 被 drop 时，锁会自动释放。
/// 这个结构体还允许你通过 `Deref` 和 `DerefMut` trait 访问受保护的数据。
pub struct LoggedMutexGuard<'a, T> {
    inner: tokio::sync::MutexGuard<'a, T>,
    mutex: &'a LoggedMutex<T>,
    holder_id: usize,
}

impl<T> Drop for LoggedMutexGuard<'_, T> {
    fn drop(&mut self) {
        self.mutex.holder_id.store(0, Ordering::Relaxed);
        debug!(
            name = %self.mutex.name,
            holder_id = self.holder_id,
            // module = module_path!(), // 添加模块名
            // line = line!(),             // 添加行号
            "Unlock"
        );
    }
}

impl<T> std::ops::Deref for LoggedMutexGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<T> std::ops::DerefMut for LoggedMutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tracing_subscriber::fmt;

    #[tokio::test]
    async fn test_lock_and_unlock() {
        // 初始化 tracing 以查看日志输出
        let subscriber = fmt::Subscriber::builder()
            .with_max_level(tracing::Level::DEBUG) // 设置日志级别为 DEBUG 或更低
            .finish();
        tracing::subscriber::set_global_default(subscriber).unwrap();

        let mutex = Arc::new(LoggedMutex::new("test_lock", 0));
        let m1 = mutex.clone();
        let m2 = mutex.clone();

        let task1 = tokio::spawn(async move {
            let guard = m1.lock().await;
            debug!(
                holder_id = m1.current_holder(),
                // module = module_path!(),
                // line = line!(),
                "Task 1 获取锁"
            );
            assert_eq!(*guard, 0);
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; // 模拟持有锁一段时间
            drop(guard);
            debug!(
                holder_id = m1.current_holder(),
                // module = module_path!(),
                // line = line!(),
                "Task 1 释放锁"
            );
        });

        let task2 = tokio::spawn(async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(5)).await; // 确保 task1 先获取锁
            let guard = m2.lock().await;
            debug!(
                holder_id = m2.current_holder(),
                // module = module_path!(),
                // line = line!(),
                "Task 2 获取锁"
            );
            assert_eq!(*guard, 0);
            drop(guard);
            debug!(
                holder_id = m2.current_holder(),
                // module = module_path!(),
                // line = line!(),
                "Task 2 释放锁"
            );
        });

        let _ = tokio::join!(task1, task2); // 等待两个任务完成
        assert_eq!(*mutex.lock().await, 0);
    }

    #[tokio::test]
    async fn test_try_lock() {
        let mutex = Arc::new(LoggedMutex::new("test_try_lock", 0));

        let guard1 = mutex.try_lock();
        assert!(guard1.is_some());
        assert_eq!(mutex.current_holder().unwrap(), 1);
        drop(guard1);
        assert!(mutex.current_holder().is_none());

        let _guard2 = mutex.lock().await;
        assert!(mutex.try_lock().is_none());
    }

    #[tokio::test]
    async fn test_holder_information() {
        let mutex = Arc::new(LoggedMutex::new("test_holder_info", 0));

        assert_eq!(mutex.current_holder(), None);

        let guard1 = mutex.lock().await;
        assert!(mutex.current_holder().is_some());
        let holder1 = mutex.current_holder().unwrap();
        drop(guard1);
        assert_eq!(mutex.current_holder(), None);

        let guard2 = mutex.lock().await;
        assert!(mutex.current_holder().is_some());
        let holder2 = mutex.current_holder().unwrap();
        assert_ne!(holder1, holder2); // 确保持有者 ID 不同
        drop(guard2);
    }

    #[tokio::test]
    async fn test_execute() {
        let mutex = Arc::new(LoggedMutex::new("test_execute", String::from("hello")));
        let result = mutex
            .execute(|s| {
                s.push_str(", world!");
                s.len()
            })
            .await;
        assert_eq!(result, 13);
        assert_eq!(*mutex.lock().await, "hello, world!");
    }

    // #[tokio::test]
    // async fn test_execute_async() {
    //     let mutex = Arc::new(LoggedMutex::new("test_execute_async", 10));
    //     let result = mutex.execute_async(|n| async {
    //         *n += 5;
    //         *n * 2
    //     }).await;
    //     assert_eq!(result, 30);
    //     assert_eq!(*mutex.lock().await, 15);
    // }

    // #[tokio::test]
    // async fn test_execute_async() {
    //     let mutex = Arc::new(LoggedMutex::new("test_execute_async", 10));
    //     let result = mutex.execute_async(|n| async move {
    //         *n += 5;
    //         *n * 2
    //     }).await;
    //     assert_eq!(result, 30);
    //     assert_eq!(*mutex.lock().await, 15);
    // }
}
